"""
海龟交易策略(pyfolio实现)
策略逻辑：
1. 使用20日突破作为入场信号
2. 使用10日低点作为止损点
3. 使用ATR计算仓位大小
"""

import pandas as pd
import pyfolio as pf
import numpy as np

def turtle_strategy(prices, entry_period=20, exit_period=10, atr_period=14, risk_pct=0.01):
    """
    海龟交易策略
    
    参数:
        prices: 价格DataFrame(需包含High,Low,Close)
        entry_period: 突破周期
        exit_period: 止损周期
        atr_period: ATR计算周期
        risk_pct: 每笔交易风险比例
    
    返回:
        positions: 持仓序列
    """
    # 计算突破通道
    high = prices['High']
    low = prices['Low']
    close = prices['Close']
    
    entry_high = high.rolling(entry_period).max()
    exit_low = low.rolling(exit_period).min()
    
    # 计算ATR
    tr = pd.concat([
        high - low,
        abs(high - close.shift(1)),
        abs(low - close.shift(1))
    ], axis=1).max(axis=1)
    atr = tr.rolling(atr_period).mean()
    
    # 生成交易信号
    positions = pd.Series(0, index=prices.index)
    in_position = False
    
    for i in range(entry_period, len(prices)):
        if not in_position and close[i] > entry_high[i-1]:
            positions[i] = (risk_pct * close[i]) / atr[i]  # 根据ATR计算仓位
            in_position = True
        elif in_position and close[i] < exit_low[i-1]:
            positions[i] = 0
            in_position = False
    
    return positions

if __name__ == '__main__':
    # 示例用法
    import yfinance as yf
    
    # 获取示例数据
    data = yf.download('AAPL', start='2020-01-01', end='2023-01-01')
    
    # 运行策略
    positions = turtle_strategy(data)
    
    # 使用pyfolio分析结果
    returns = positions.shift(1) * data['Close'].pct_change()
    pf.create_returns_tear_sheet(returns)